import sys
import requests
import threading
import os
import time
from PyQt5.QtWidgets import QApplication, QWidget, QVBoxLayout, QHBoxLayout, QLabel, QLineEdit, QPushButton, QProgressBar
from PyQt5.QtCore import QObject, pyqtSignal

# 下载文件的指定部分
def download_part(url, start, end, filename, part_num, progress_signal, speed_signal, size_signal, pause_event):
    part_filename = f'{filename}.part{part_num}'
    if os.path.exists(part_filename):
        downloaded = os.path.getsize(part_filename)
        start += downloaded
    else:
        downloaded = 0

    headers = {'Range': f'bytes={start}-{end}'}
    try:
        response = requests.get(url, headers=headers, stream=True)
        total_size = end - start + 1
        last_time = time.time()
        last_downloaded = downloaded
        with open(part_filename, 'ab') as f:
            for chunk in response.iter_content(chunk_size=1024):
                if pause_event.is_set():
                    while pause_event.is_set():
                        time.sleep(0.1)
                if chunk:
                    f.write(chunk)
                    downloaded += len(chunk)
                    progress = (downloaded / total_size) * 100
                    progress_signal.emit(part_num, progress)

                    # 计算下载速度
                    current_time = time.time()
                    elapsed_time = current_time - last_time
                    if elapsed_time > 0:
                        speed = (downloaded - last_downloaded) / elapsed_time
                        speed_signal.emit(speed)
                        last_time = current_time
                        last_downloaded = downloaded

                    # 发送已下载大小信号
                    size_signal.emit(downloaded)
    except Exception as e:
        print(f"Error downloading part {part_num}: {e}")


# 合并所有部分文件
def merge_parts(filename, num_parts):
    with open(filename, 'wb') as outfile:
        for i in range(num_parts):
            part_filename = f'{filename}.part{i}'
            with open(part_filename, 'rb') as infile:
                outfile.write(infile.read())
            os.remove(part_filename)


# 自定义信号类
class DownloadProgressSignal(QObject):
    progress_updated = pyqtSignal(int, float)
    speed_updated = pyqtSignal(float)
    size_updated = pyqtSignal(int)
    download_finished = pyqtSignal()


# 主下载函数
def download_file(url, filename, num_threads=4, progress_signal=None, speed_signal=None, size_signal=None, pause_event=None, finished_signal=None):
    try:
        response = requests.head(url)
        file_size = int(response.headers.get('content-length', 0))
        if file_size == 0:
            print("Unable to get file size.")
            return

        part_size = file_size // num_threads
        threads = []
        for i in range(num_threads):
            start = i * part_size
            end = start + part_size - 1 if i < num_threads - 1 else file_size - 1
            thread = threading.Thread(target=download_part, args=(url, start, end, filename, i, progress_signal, speed_signal, size_signal, pause_event))
            threads.append(thread)
            thread.start()

        for thread in threads:
            thread.join()

        merge_parts(filename, num_threads)
        print(f"Download completed: {filename}")
        if finished_signal:
            finished_signal.emit()
    except Exception as e:
        print(f"Error downloading file: {e}")
        if finished_signal:
            finished_signal.emit()


class DownloadApp(QWidget):
    def __init__(self):
        super().__init__()
        self.pause_event = threading.Event()
        self.total_downloaded = 0
        self.initUI()

    def initUI(self):
        # 创建布局
        main_layout = QVBoxLayout()

        # 下载链接输入框
        url_layout = QHBoxLayout()
        url_label = QLabel("下载链接:")
        self.url_input = QLineEdit()
        self.url_input.setText("https://www.python.org/ftp/python/3.12.9/python-3.12.9-amd64.exe")
        url_layout.addWidget(url_label)
        url_layout.addWidget(self.url_input)
        main_layout.addLayout(url_layout)

        # 保存路径输入框
        path_layout = QHBoxLayout()
        path_label = QLabel("保存路径:")
        self.path_input = QLineEdit()
        self.path_input.setText("D:/download/python-3.12.9-amd64.exe")
        path_layout.addWidget(path_label)
        path_layout.addWidget(self.path_input)
        main_layout.addLayout(path_layout)

        # 开始下载按钮
        self.start_button = QPushButton("开始下载")
        self.start_button.clicked.connect(self.start_download)
        main_layout.addWidget(self.start_button)

        # 暂停/继续按钮
        self.pause_button = QPushButton("暂停下载")
        self.pause_button.clicked.connect(self.pause_or_resume_download)
        self.pause_button.setEnabled(False)
        main_layout.addWidget(self.pause_button)

        # 进度条
        self.progress_bars = []
        for i in range(4):
            progress_bar = QProgressBar()
            progress_bar.setRange(0, 100)
            main_layout.addWidget(progress_bar)
            self.progress_bars.append(progress_bar)

        # 下载速度显示
        self.speed_label = QLabel("下载速度: 0.00 KB/s")
        main_layout.addWidget(self.speed_label)

        # 已下载文件大小显示
        self.size_label = QLabel("已下载大小: 0.00 KB")
        main_layout.addWidget(self.size_label)

        self.setLayout(main_layout)
        self.setWindowTitle('下载工具')
        self.setGeometry(300, 300, 400, 300)
        self.show()

    def start_download(self):
        url = self.url_input.text()
        filename = self.path_input.text()
        if not url or not filename:
            print("请输入下载链接和保存路径。")
            return

        self.signal = DownloadProgressSignal()
        self.signal.progress_updated.connect(self.update_progress)
        self.signal.speed_updated.connect(self.update_speed)
        self.signal.size_updated.connect(self.update_size)
        self.signal.download_finished.connect(self.reset_download_state)

        self.download_thread = threading.Thread(target=download_file,
                                                args=(url, filename, 4, self.signal.progress_updated,
                                                      self.signal.speed_updated, self.signal.size_updated,
                                                      self.pause_event, self.signal.download_finished))
        self.download_thread.start()
        self.start_button.setEnabled(False)
        self.pause_button.setEnabled(True)

    def pause_or_resume_download(self):
        if self.pause_event.is_set():
            self.pause_event.clear()
            self.pause_button.setText("暂停下载")
        else:
            self.pause_event.set()
            self.pause_button.setText("继续下载")

    def update_progress(self, part_num, progress):
        self.progress_bars[part_num].setValue(int(progress))

    def update_speed(self, speed):
        # 将速度转换为 KB/s
        speed_kb = speed / 1024
        self.speed_label.setText(f"下载速度: {speed_kb:.2f} KB/s")

    def update_size(self, size):
        self.total_downloaded = size
        # 将字节转换为合适的单位
        if self.total_downloaded < 1024:
            size_str = f"{self.total_downloaded:.2f} B"
        elif self.total_downloaded < 1024 ** 2:
            size_str = f"{self.total_downloaded / 1024:.2f} KB"
        elif self.total_downloaded < 1024 ** 3:
            size_str = f"{self.total_downloaded / (1024 ** 2):.2f} MB"
        else:
            size_str = f"{self.total_downloaded / (1024 ** 3):.2f} GB"
        self.size_label.setText(f"已下载大小: {size_str}")

    def reset_download_state(self):
        # 重置进度条
        # for progress_bar in self.progress_bars:
        #     progress_bar.setValue(0)
        # 重置下载速度显示
        # self.speed_label.setText("下载速度: 0.00 KB/s")
        # 重置已下载大小显示
        # self.size_label.setText("已下载大小: 0.00 KB")
        # 重置按钮状态
        self.start_button.setEnabled(True)
        self.pause_button.setEnabled(False)
        self.pause_button.setText("暂停下载")
        # 重置暂停事件
        self.pause_event.clear()
        # 重置已下载字节数
        # self.total_downloaded = 0


if __name__ == '__main__':
    app = QApplication(sys.argv)
    ex = DownloadApp()
    sys.exit(app.exec_())